# -*- coding: utf-8 -*-

'''
Created on 26/11/2011

@author: Luis M. Ugalde
www.forondarena.net

Copyright © 2011 Luis M. Ugalde
Licensed under the terms of the MIT License
(see LICENSE file for details)

This software can't be considered to be even in Alpha state. It still needs a lot of work.
'''


from . import abstractstatestore
import shelve
import zmq
import pickle
from statestores.statestoreexception import StateStoreException


class StateStoreShelve(abstractstatestore.AbstractStateStore):
    '''
    Class to keep state for multiline records.
    This state store will know when to end with a special END message sent from the parser.
    This class uses shelve to keep state.
    '''


    def __init__(self, shelvefilename = "", shelvewriteback = True, *args, **kwargs):
        '''
        Define the shelve database that will be used to keep state of records.
        '''
        
        # I'm not moving the shelve related methods to the abstract class, as this is a Shelve-only thing.

        self.filename = shelvefilename
        self.writeback = shelvewriteback
        print "{}: StateStore instance created with {} shelve filename".format(self.__class__.__name__, self.filename)
        try:
            self.db = shelve.open(self.filename, writeback = self.writeback)
        except Exception:
            tmperror = "Something went wrong with the shelve " + self.filename
            raise StateStoreException(tmperror)

        abstractstatestore.AbstractStateStore.__init__(self, *args, **kwargs)

        
    def consumeLine(self, line):
        '''
        '''
        
        try:
            data = pickle.loads(line)
        except pickle.PickleError:
            pass
        if __debug__:
            print "{}-{}: Processing line: {}".format(self.__class__.__name__, self.statestorebindaddress, data)

        
        if len(data) != 4:
            # This should be handled somehow, but it's probably not enough to terminate
            print "{}-{}: This record is invalid. It must have more than 4 elements. {}".format(self.__class__.__name__, self.statestorebindaddress, data)
            return

        # Let's give names to data records
        recordtype = data[0]
        recordmodifier = data[1]
        recordID = data[2]
        recordpayload = data[3]
        
        # I'll leave pickle and zeromq exception handling to the caller for the moment.

        if recordtype == 'NODATA':
            pickledbuffer = pickle.dumps('NODATA')
            self.statestorebindsocket.send(pickledbuffer)
            if __debug__:
                print "{}-{}: Sent NODATA record.".format(self.__class__.__name__, self.statestorebindaddress)
            return

        if recordtype == 'SINGLELINE':
            pickledbuffer = pickle.dumps(recordpayload)
            self.statestorebindsocket.send(pickledbuffer)
            if __debug__:
                print "{}-{}: Sent SINGLELINE record.".format(self.__class__.__name__, self.statestorebindaddress)
            return

        elif recordtype == 'KEYDATA':
            if recordmodifier == 'NEWID':
                try:
                    if self.db[recordID]:
                        print "{}-{}: Duplicate key. Removing the old one and keeping this: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                        del self.db[recordID]
                        self.db[recordID] = recordpayload
                except KeyError:
                    self.db[recordID] = recordpayload
                    if __debug__:
                        print "{}-{}: Storing new key: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
            elif recordmodifier == 'DATA':
                try:
                    for key in recordpayload.keys():
                        self.db[recordID][key] = recordpayload[key]
                except KeyError:
                    if __debug__:
                        print "{}-{}: The ID does not exist. Cant add data: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
            elif recordmodifier == 'ENDID':
                try:
                    for key in recordpayload.keys():
                        self.db[recordID][key] = recordpayload[key]
                    if __debug__:
                        print "{}-{}: Ending message: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                    pickledbuffer = pickle.dumps(self.db[recordID])
                    self.statestorebindsocket.send(pickledbuffer)
                    del self.db[recordID]
                except KeyError:
                    if __debug__:
                        print "{}-{}: The ID does not exist. Cant end record: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
        
        elif recordtype == 'KEYDATAMULTIPLE':
            try:
                tmpremote = self.db[recordID]
            except KeyError:
                if __debug__:
                    print "{}-{}: The ID does not exist. Cant add keydatamultiple: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
            else:
                # I should check whether the received data is valid
                try:
                    tmparray = tmpremote[recordmodifier]
                    tmparray.append(recordpayload)
                except KeyError:
                    tmparray = [recordpayload]
                    tmpremote[recordmodifier] = tmparray
                    self.db[recordID] = tmpremote
                    
        # This is too qmail specific. I need more samples to generalize
        elif recordtype == 'SUBKEYDATA':
            if recordmodifier == 'NEWID':
                try:
                    if __debug__:
                        print "{}-{}: Storing a new subkey for key {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                        print "{}-{}: The subkey id is {}.".format(self.__class__.__name__, self.statestorebindaddress, '_subqueueid_' + recordpayload['subqueueid'])
                    tmpremote = self.db[recordID]
                    try:
                        tmpdict = self.db['_subqueueid_' + recordpayload['subqueueid']]
                        if tmpdict['queueid'] != recordID or tmpdict['rcptto'] != recordpayload['rcptto']:
                            if __debug__:
                                print "{}-{}: The new subkey exists for a different queuid or rcptto. Deleting entry {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                            del self.db['_subqueueid_' + recordpayload['subqueueid']]
                            return
                    except KeyError:
                        if __debug__:
                            print "{}-{}: The subkey does not exist. Good. Creating it: {}.".format(self.__class__.__name__, self.statestorebindaddress, '_subqueueid_' + recordpayload['subqueueid'])
                        try:
                            self.db['_subqueueid_' + recordpayload['subqueueid']] = {'queueid': recordID, 'rcptto':  recordpayload['rcptto']}
                        except KeyError:
                            if __debug__:
                                print "{}-{}: The new subkey data is bogus, discarding entry for queueid {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                            return
                except KeyError:
                    if __debug__:
                        print "{}-{}: Cant add a subkey for a non existing key {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                    return
                
            elif recordmodifier == 'DATA':
                try:
                    if __debug__:
                        print "{}-{}: Adding subkey data to record: {}.".format(self.__class__.__name__, self.statestorebindaddress, '_subqueueid_' + recordID)
                    tmpdict = self.db['_subqueueid_' + recordID]
                    tmpdict.update(recordpayload)
                    self.db['_subqueueid_' + recordID] = tmpdict
                except KeyError:
                    if __debug__:
                        print "{}-{}: The subkey ID does not exist. Cant add data: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                    return

            elif recordmodifier == 'ENDID':
                try:
                    if __debug__:
                        print "{}-{}: Ending message: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                        print "{}-{}: Fetching from database: {}.".format(self.__class__.__name__, self.statestorebindaddress, '_subqueueid_' + recordID)
                    tmpdict = self.db['_subqueueid_' + recordID]
                    tmpqueueid = tmpdict['queueid']
                    del tmpdict['queueid']
                    tmprecord = self.db[tmpqueueid]
                except KeyError:
                    if __debug__:
                        print "{}-{}: The subkey ID or the key ID do not exist. Cant end record: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                else:
                    try:
                        # I use rcptto, but this should also be read from the pickled data
                        tmprecord['rcptto'].append(tmpdict)
                        if __debug__:
                            print "{}-{}: Appending to array for id {} with contents: {}.".format(self.__class__.__name__, self.statestorebindaddress, tmpqueueid, tmpdict)
                    except KeyError:
                        tmprecord['rcptto'] = [tmpdict]
                        if __debug__:
                            print "{}-{}: Creating a new array for id {} with contents: {}.".format(self.__class__.__name__, self.statestorebindaddress, tmpqueueid, tmpdict)
                    self.db[tmpqueueid] = tmprecord
                    if __debug__:
                        print "{}-{}: Delete finished record: {}.".format(self.__class__.__name__, self.statestorebindaddress, recordID)
                    del self.db['_subqueueid_' + recordID]


    def processLines(self):
        '''
        Process lines that the parser sends.
        '''
        
        mustexit = 0
        idlesockets = 0
        
        while True:
            socks = dict(self.poller.poll(3000))
            if socks:
                idlesockets = 0
                for i in self.parserconnectsockets:
                    if i in socks and socks.get(i) == zmq.POLLIN:
                        try:
                            line = i.recv()
                            self.consumeLine(line)
                        except (pickle.PickleError, zmq.ZMQError):
                            # Not sure if I should end here, but I'll do until I find another way
                            # It would probably be nice to propagate an ERROR message
                            print "{}-{}: Error with pickle or ZeroMQ. Ending {}.".format(self.__class__.__name__, self.statestorebindnotifyaddress, line)
                            break 
                if self.parserconnectnotifysocket in socks and socks.get(self.parserconnectnotifysocket) == zmq.POLLIN:
                    line = self.parserconnectnotifysocket.recv()
                    print "{}-{}: The line is {}.".format(self.__class__.__name__, self.statestorebindnotifyaddress, line)
                    if line == 'END':
                        if __debug__:
                            print "{}-{}: Received END signal. Setting END flag.".format(self.__class__.__name__, self.statestorebindnotifyaddress)
                        mustexit = 1
            else:
                if mustexit > 0:
                    print "{}-{}: No data left at State Store for 3 seconds wind end signal. Ending.".format(self.__class__.__name__, self.statestorebindaddress)
                    self.statestorebindnotifysocket.send('END')
                    print "{}-{}: Sent END signal. Trying to end State Store".format(self.__class__.__name__, self.statestorebindnotifyaddress)
                    break
                if idlesockets > 3:
                    print "{}-{}: Too long with no activity from sockets. Should probably do something".format(self.__class__.__name__, self.statestorebindaddress)
                    idlesockets = 0
                idlesockets += 1

        self.closeSources()
        self.closeNotifySources()
        self.closeBindStateStore()
        self.closeBindNotifyStateStore()
        self.closeContext()
        self.db.close()


if __name__ == '__main__':
    print "Please, do not run this"
